package p.ithorns.integration.zk.id;

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.atomic.AtomicValue;
import org.apache.curator.framework.recipes.atomic.DistributedAtomicLong;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.springframework.util.StringUtils;

import javax.annotation.PostConstruct;
import java.util.Map;
import java.util.concurrent.*;

/**
 * ZkIdGenerator
 *
 * @author ithorns
 * @version 1.0
 * @date 2021/9/7 20:58
 */
@Slf4j
public class ZkIdGenerator implements IdGenerator, Watcher {

    private final ZkProperties zkProperties;

    private CuratorFramework client;

    private ExecutorService executor;

    public ZkIdGenerator(ZkProperties properties) throws Exception {
        this.zkProperties = properties;
        init();
    }

    @Override
    @PostConstruct
    public void init() throws Exception {
        client = CuratorFrameworkFactory.builder()
                .connectString(zkProperties.getHosts())
                .sessionTimeoutMs(zkProperties.getSessionTimeout())
                .connectionTimeoutMs(zkProperties.getConnTimeout())
                .retryPolicy(new ExponentialBackoffRetry(zkProperties.getRetrySleep(), zkProperties.getRetryTimes()))
                .namespace("id_gen")
                .build();

        ThreadFactory namedThreadFactory = new ThreadFactoryBuilder()
                .setNameFormat("id-pool-%d").build();
        executor = new ThreadPoolExecutor(
                2, 4,
                0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<>(1024),
                namedThreadFactory,
                new ThreadPoolExecutor.AbortPolicy());

        client.start();

        // 判断父节点/root节点是否存在
//        String root = zkProperties.getRoot();
//        Stat stat = client.checkExists().forPath(root);
//        if (stat == null) {
//            client.create()
//                    .withMode(CreateMode.PERSISTENT)
//                    .forPath(root, null);
//        }

//        final NodeCache nodeCache = new NodeCache(client,  root, true);
//        nodeCache.start();
//        nodeCache.getListenable().addListener(() -> {
//            String path = nodeCache.getCurrentData().getPath();
//            log.info("事件监听result={}", path);
//        });
    }

    // 序列化集合
    private final Map<String, DistributedAtomicLong> zkSequence = Maps.newConcurrentMap();

    private final RetryPolicy retryPolicy = new ExponentialBackoffRetry(200, 10);

    public Long getNextId(String name) {
        try {
            DistributedAtomicLong distAtomicLong = this.zkSequence.computeIfAbsent(name,
                    key -> new DistributedAtomicLong(client,  name, retryPolicy));
            AtomicValue<Long> sequence = distAtomicLong.increment();
            if (sequence.succeeded()) {
                return sequence.postValue();
            } else {
                return null;
            }
        } catch (Exception e) {
            log.error("获取[{}]Sequence错误:{}", name, e);
        }
        return null;
    }

    @Override
    public String getNextId(IdBean idBean) {
        final String path = idBean.getPath();
        String id = null;
        try {
            id = client.create()
                    .creatingParentsIfNeeded()
                    .withMode(CreateMode.EPHEMERAL_SEQUENTIAL)
                    .forPath(path);
        } catch (Exception e) {
            throw new IDGenerateException("Failed to get next ID: create EPHEMERAL_SEQUENTIAL node error.", e);
        }

        if (!StringUtils.hasText(id)) {
            throw new IDGenerateException("Get next ID: null.");
        }
        // 删除创建的持久化顺序节点
        final String finalId = id;
        executor.execute(() -> {
            try {
                client.delete().forPath(finalId);
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        });

        return split(id);
    }

    @Override
    public void destroy() {
        executor.shutdown();
        this.client.delete();
        this.client.close();
    }

    @Override
    public void process(WatchedEvent watchedEvent) {
        try {
            //EventType = None时
            if (watchedEvent.getType() == Watcher.Event.EventType.None) {
                if (watchedEvent.getState() == Watcher.Event.KeeperState.SyncConnected) {
                    log.info("连接成功");
                } else if (watchedEvent.getState() == Watcher.Event.KeeperState.Disconnected) {
                    log.info("断开连接");
                } else if (watchedEvent.getState() == Watcher.Event.KeeperState.Expired) {
                    log.info("会话超时");
                    // 超时后服务器端已经将连接释放，需要重新连接服务器端
                } else if (watchedEvent.getState() == Watcher.Event.KeeperState.AuthFailed) {
                    log.info("认证失败");
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private String split(String path) {
        int index = path.lastIndexOf(IdBean.ID_SEPARATOR);
        if (index >= 0) {
            String intStr = path.substring(index + 1);
            // return Integer.parseInt(intStr) + 1;
            return StringUtils.trimAllWhitespace(intStr)
                    .replaceFirst("^(0*)", "");
        }

        throw new IDGenerateException("An exception occurred during ID generation.");
    }

}
